#include "builder.hh"
#include "tree.hh"
#include "node.hh"
#include "constants.hh"
#include "cppn/cppn.hh"
#include <iostream>
#include <algorithm>
#include <iterator>
#include <cassert>
using namespace std;

namespace brain {

    void
    Builder::construct () {
        if (hidden.size() > 0) return;

        set<Node> connected_nodes;
        set<Node> forward_nodes;
        set<Node> backward_nodes;

        set<Node> unexplored_forward_nodes { inputs };
        set<Node> unexplored_backward_nodes { outputs };

        for (int i = 0; i < NETWORK_ITERATION_LEVEL; ++i) {
            set<Node> new_forward_nodes;
            set<Node> new_backward_nodes;
            cout << "There are " << unexplored_forward_nodes.size() << " unexplored forward nodes!" << endl;
            cout << "There are " << unexplored_backward_nodes.size() << " unexplored backward nodes!" << endl;

            for (Node const& src: unexplored_forward_nodes) {
                set<Connection> const& new_connections = tree->get_connections_for(src, true);
                for (Connection const& c: new_connections) {
                    Node const& dst = c.get_dst();
                 
                   if (inputs.count(dst) == 0) {
                        if (connected_nodes.count(dst) == 1) {
                            connected_nodes.insert(src);
                            forward_nodes.erase(src);

                        } else if (backward_nodes.count(dst) == 1 || outputs.count(dst) == 1) {
                            connected_nodes.insert(dst);
                            connected_nodes.insert(src);
                            forward_nodes.erase(dst);
                            forward_nodes.erase(src);
                            backward_nodes.erase(dst);

                        } else if (forward_nodes.count(dst) == 0) {
                            new_forward_nodes.insert(dst);
                            forward_nodes.insert(dst);
                        }
                    }
                }
            }

            for (Node const& dst: unexplored_backward_nodes) {
                set<Connection> const& new_connections = tree->get_connections_for(dst, false);
                for (Connection const& c: new_connections) {
                    Node const& src = c.get_src();

                    if (outputs.count(src) == 0) {
                        if (connected_nodes.count(src) == 1) {
                            connected_nodes.insert(dst);
                            backward_nodes.erase(dst);

                        } else if (forward_nodes.count(src) == 1 || inputs.count(dst) == 1) {
                            connected_nodes.insert(src);
                            connected_nodes.insert(dst);
                            backward_nodes.erase(src);
                            backward_nodes.erase(dst);
                            forward_nodes.erase(src);

                        } else if (backward_nodes.count(src) == 0) {
                            new_backward_nodes.insert(src);
                            backward_nodes.insert(src);
                        }
                    }
                }
            }

            swap(unexplored_forward_nodes, new_forward_nodes);
            swap(unexplored_backward_nodes, new_backward_nodes);
        }


        for (int i = 0; i < NETWORK_ITERATION_LEVEL; ++i) {
            int forward_plus = 0;
            int backward_plus = 0;

            set<Node> backtrace_forward_nodes { forward_nodes };
            for (Node const& node: backtrace_forward_nodes) {
                set<Connection> const& new_connections = tree->get_connections_for(node, true);
                for (Connection const& c: new_connections) {
                    if (inputs.count(c.get_dst()) == 0) {
                        if (connected_nodes.count(c.get_dst()) == 1) {
                            connected_nodes.insert(node);
                            forward_nodes.erase(node);
                            forward_plus++;
                        }
                    }
                }
            }

            set<Node> backtrace_backward_nodes { backward_nodes };
            for (Node const& node: backtrace_backward_nodes) {
                set<Connection> const& new_connections = tree->get_connections_for(node, false);
                for (Connection const& c: new_connections) {
                    if (outputs.count(c.get_src()) == 0) {
                        if (connected_nodes.count(c.get_src()) == 1) {
                            connected_nodes.insert(node);
                            backward_nodes.erase(node);
                            connections.insert(c);
                            backward_plus++;
                        }
                    }
                }
            }

            cout << "Additional forward connections: " << forward_plus << endl;
            cout << "Additional backward connections: " << backward_plus << endl;

        }

        cout << "Connected nodes: " << connected_nodes.size() << endl;
        cout << "Nodes connected to input, but not to output: " << forward_nodes.size() << endl;
        cout << "Nodes connected to output, but not to input: " << backward_nodes.size() << endl;


        copy_if(connected_nodes.begin(), connected_nodes.end(), inserter(hidden, hidden.begin()),
                [&] (Node const& n) {
                    return inputs.count(n) == 0 && outputs.count(n) == 0;
                });

        copy(inputs.begin(), inputs.end(), inserter(connected_nodes, connected_nodes.begin()));
        copy(outputs.begin(), outputs.end(), inserter(connected_nodes, connected_nodes.begin()));

        for (Node const& n: connected_nodes) {
            for (Connection const& c: tree->get_existing_connections_for(n)) {
                if ((n == c.get_src() && connected_nodes.count(c.get_dst()) == 1)
                    || (n == c.get_dst() && connected_nodes.count(c.get_src()) == 1)) {
                    
                    if (inputs.count(c.get_dst()) == 0 && outputs.count(c.get_src()) == 0) {
                        connections.insert(c);
                    }
                }
            }
        }

        cout << "connections: " << connections.size() << endl;
    }

    void
    Builder::print() {
        for (auto const& node: inputs) {
                print(node);
        }
        for (auto const& node: hidden) {
                print(node);  
        }
    }

    void
    Builder::print(Node const& n) {
        for (Connection const& c: connections) {
            Node dst = c.get_dst();

            if (c.get_src() == n) {
                if (hidden.count(dst) == 1)
                    cout << n << " -> " << *hidden.find(dst) << '\n';
                else if (outputs.count(c.get_dst()) == 1)
                    cout << n << " -> " << *outputs.find(dst) << " (output)\n";                
            }
        }
    }

}








